package 中等.搜索.深度优先搜索;

import java.util.*;

/**
 * 现有一棵由 n 个节点组成的无向树，节点编号从 0 到 n - 1 ，共有 n - 1 条边。
 * 给你一个二维整数数组 edges ，长度为 n - 1 ，其中 edges[i] = [ai, bi] 表
 * 示树中节点 ai 和 bi 之间存在一条边。另给你一个整数数组 restricted 表示 受限 节点。
 * 在不访问受限节点的前提下，返回你可以从节点 0 到达的 最多 节点数目。
 * 注意，节点 0 不 会标记为受限节点。
 *
 * @ https://leetcode.cn/problems/reachable-nodes-with-restrictions/description/?envType=daily-question&envId=2024-03-02
 * @date 2024/03/02
 */
public class 受限条件下可到达节点的数目_2368 {

    public static void main(String[] args) {

    }

    /**
     * 构造图+深度优先搜索
     */
    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        // 节点 key 可以到达的下一个节点有哪些
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int[] edge : edges) {
            List<Integer> list = map.get(edge[0]);
            if (list == null) {
                list = new ArrayList<>();
                map.put(edge[0], list);
            }
            list.add(edge[1]);
            List<Integer> list1 = map.get(edge[1]);
            if (list1 == null) {
                list1 = new ArrayList<>();
                map.put(edge[1], list1);
            }
            list1.add(edge[0]);
        }

        Set<Integer> restrictedSet = new HashSet<>();
        for (int node : restricted) {
            restrictedSet.add(node);
        }
        return dfs(0, map, new boolean[n], restrictedSet);
    }

    private int dfs(int node, Map<Integer, List<Integer>> map, boolean[] visited, Set<Integer> restrictedSet) {
        if (visited[node] || restrictedSet.contains(node)) {
            return 0;
        }
        visited[node] = true;
        int ans = 1;

        List<Integer> nextList = map.get(node);
        if (nextList != null) {
            for (Integer nextNode : nextList) {
                ans += dfs(nextNode, map, visited, restrictedSet);
            }
        }
        return ans;
    }

}
